# Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
import argparse
import secrets
import os
import shutil
import getpass
import json
import zipfile
import requests
import secrets
import threading
import time
from tqdm import tqdm
import queue
import time
def create_zip(source_dir, zip_file_path):

    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(source_dir):
            for file in files:
                file_path = os.path.join(root, file)
                zipf.write(file_path, os.path.relpath(file_path, source_dir))


def send_zip_with_form_data(response_queue, user_sub, user_passwd, method, url, zip_file_path):
    while 1:
        zip_file=open(zip_file_path, 'rb')
        files = {'zip_file': (os.path.basename(zip_file_path), zip_file)}
        data = {
            'user_sub': user_sub,
            'user_passwd': user_passwd,
            'method': method
        }
        response = requests.post(url, files=files, data=data)
        if response.json()['code']==201:
            time.sleep(5)
            zip_file.close()
            continue
        else:
            zip_file.close()
            break
    
    response_queue.put(response)


def send_corpus_name(user_sub, user_passwd, method, url, corpus_name):
    data = {
        'user_sub': user_sub,
        'user_passwd': user_passwd,
        'corpus_name': corpus_name,
        'method': method
    }
    response = requests.post(url, data=data)
    return response


def work(args):
    choice = args['method']
    user_sub = args['user_sub']
    user_passwd = args['user_passwd']
    corpus_name = args['corpus_name']
    corpus_dir = args['corpus_dir']
    ssl_enable = args['ssl_enable']
    host = args['host']
    port = args['port']
    if ssl_enable:
        url = 'https://'+host+':'+port+'/corpus'
    else:
        url = 'http://'+host+':'+port+'/corpus'
    if choice == "init_user_info":
        if user_sub is None:
            print("请输入你的用户名")
            exit(1)
        if user_passwd is None:
            print("请输入你的用户名")
            exit(1)
        if os.path.exists('user_info.json'):
            os.remove('user_info.json')
        user_info = {}
        user_info['user_sub'] = user_sub
        user_info['user_passwd'] = user_passwd
        with open('user_info.json', 'w') as f:
            json.dump(user_info, f)
    elif choice == "up_corpus":
        user_info = {}
        try:
            with open('user_info.json', 'r') as f:
                user_info = json.load(f)
        except:
            pass
        if user_sub is None:
            user_sub = user_info.get('user_sub', '')
        if user_passwd is None:
            user_passwd = user_info.get('user_passwd', '')
        zip_file_path = os.path.join('./' + secrets.token_hex(16)+'.zip')
        create_zip(corpus_dir, zip_file_path)
        response_queue = queue.Queue()
        upload_thread = threading.Thread(target=send_zip_with_form_data, args=(
            response_queue, user_sub, user_passwd, choice, url, zip_file_path))
        upload_thread.start()
        with tqdm(desc="等待上传中（会有1-10分钟的延迟，请稍安勿躁）...", unit="秒", dynamic_ncols=True) as pbar:
            while upload_thread.is_alive():
                pbar.update(1)
                time.sleep(1)

        upload_thread.join()

        response = response_queue.get()
        if response.json()['code'] == 200:
            print('上传语料成功')
        else:
            print('上传语料失败')
            print(response.json()['message'])
        os.remove(zip_file_path)
    elif choice == "del_corpus":
        user_info = {}
        try:
            with open('user_info.json', 'r') as f:
                user_info = json.load(f)
        except:
            pass
        if user_sub is None:
            user_sub = user_info.get('user_sub', '')
        if user_passwd is None:
            user_passwd = user_info.get('user_passwd', '')
        try:
            response = send_corpus_name(user_sub, user_passwd, choice, url, corpus_name)
            if response.json()['code'] == 200:
                print('删除语料成功')
            else:
                print('删除语料失败')
                print(response.json()['message'])
        except:
            print('删除语料失败')
    elif choice == "query_corpus":
        user_info = {}
        try:
            with open('user_info.json', 'r') as f:
                user_info = json.load(f)
        except:
            pass
        if user_sub is None:
            user_sub = user_info.get('user_sub', '')
        if user_passwd is None:
            user_passwd = user_info.get('user_passwd', '')
        try:
            response = send_corpus_name(user_sub, user_passwd, choice, url, corpus_name)
            if response.json()['code'] == 200:
                print('查询语料成功')
                corpus_list = response.json()['message']
                if len(corpus_list) == 0:
                    print('未查询到相关语料')
                else:
                    for corpus_name, t in corpus_list:
                        print('语料名 ', corpus_name, '上传时间 ', t)
            else:
                print('查询语料失败')
                print(response.json()['message'])
        except:
           print('查询语料失败')
    else:
        print("无效的选择")
        exit(1)


def init_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--method", type=str, required=True, choices=[
                                                                      'init_user_info',
                                                                      'up_corpus', 
                                                                      'del_corpus', 
                                                                      'query_corpus'
                                                                      ], help='''
                                                                     脚本使用模式，有初
                                                                     始化用户信息、上传
                                                                     语料、删除语料和查
                                                                     询语料功能
                                                                     ''')
    parser.add_argument("--user_sub", default=None,type=str, required=False, help="用户名")
    parser.add_argument("--user_passwd", default=None,type=str, required=False, help="用户密码")
    parser.add_argument("--corpus_name", default='',type=str, required=False, help="语料名")
    parser.add_argument("--corpus_dir", type=str, default='./docs', required=False, help="语料路径")
    parser.add_argument("--ssl_enable", type=bool, default=False, required=False, help="服务进程是否支撑https访问")
    parser.add_argument("--host", type=str, default='0.0.0.0', required=False, help="服务进程ip")
    parser.add_argument("--port", type=str, default='9005', required=False, help="服务进程端口")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = init_args()
    work(vars(args))
